import redis
from src.config import *


class RedisClient(object):
    def __init__(self, host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD):
        """
        初始化Redis连接
        :param host: 地址
        :param port: 端口
        :param password: 密码
        """
        self.db = redis.StrictRedis(host=host, port=port, password=password, decode_responses=True)

    def set(self, city, data):
        return self.db.sadd(city, data)

    def list_all(self, city):
        return self.db.smembers(city)

    def save_link(self, link):
        return self.db.sadd("links", link)

    def get_link(self):
        return self.db.spop("links")

    def save_map(self, name, mapping):
        if self.db.hlen(name) > 0:
            for key in self.db.hkeys(name):
                self.db.hdel(name, key)
        return self.db.hmset(name, mapping)

    def get_map(self, name, key):
        return self.db.hget(name, key)
